﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using RayDen.Library.Components.Surface;
using RayDen.Library.Core;
using RayDen.Library.Core.Primitives;
using RayDen.RayEngine.Core.Interface;
using RayDen.RayEngine.Core.Types;
using RayDen.RayEngine.Engines.PathTracer;
using RayDen.RayEngine.Samplers;

namespace RayDen.RayEngine.Engines.MCMC
{
    public enum PssmltSamplerState
    {
        WarmUp,
        Initialize,
        SmallMutation,
        LargeMutation,
    }

    public class PrimarySampleSpacePathSampler : PathSamplerBase
    {
        private const int SamplesPerVertex = 8;
        internal ContributionMap contributionMap;
        private RayEngineScene scene;
        private ShadowRayInfo[] secRays;
        internal RgbSpectrum Throughput;
        private int depth, mutationsCount,  tracedShadowRayCount;
        private bool specularBounce, mutate;

        private Sample XSample;
        private RgbSpectrum XContribution;

        private PssmltSamplerState SamplerState;
        private PathTracerPathState PathState;

        SurfaceIntersectionData hitInfo = null;

        public PrimarySampleSpacePathSampler()
        {
            SamplerState = PssmltSamplerState.WarmUp;
        }

        public override void InitPath(IPathProcessor buffer)
        {
            base.InitPath(buffer);
            this.scene = pathIntegrator.Scene;
            this.Radiance = new RgbSpectrum(0f);
            this.Throughput = new RgbSpectrum(1f);
            if (secRays == null)
                this.secRays = new ShadowRayInfo[scene.ShadowRayCount];
            if (pathIntegrator.Sampler.GetPass() > PssMltEngineController.WarmupIterations && SamplerState == PssmltSamplerState.WarmUp)
            {
                SamplerState = PssmltSamplerState.Initialize;
            }
        @restart:
            switch (SamplerState)
            {
                case PssmltSamplerState.WarmUp:
                    this.Sample = pathIntegrator.Sampler.GetSample();
                    this.Sample.InitSample(SamplesPerVertex * scene.MaxPathDepth);
                    this.contributionMap[this.Sample].Sample = this.Sample;
                    this.contributionMap[this.Sample].Weight = 1.0f;
                    break;
                case PssmltSamplerState.Initialize:
                    var s = pathIntegrator.Sampler.GetSample();
                    if (this.contributionMap[s].Sample != null)
                    {
                        this.Sample = this.contributionMap[s].Sample;
                    }
                    else
                    {
                        this.Sample.InitSample(SamplesPerVertex * scene.MaxPathDepth);
                        this.contributionMap[this.Sample].Sample = this.Sample;
                        this.contributionMap[this.Sample].Weight = 1.0f;
                    }
                    break;
                case PssmltSamplerState.LargeMutation:
                    s = pathIntegrator.Sampler.GetSample();
                    if (this.contributionMap[s].Sample != null)
                    {
                        this.Sample = this.contributionMap[s].Sample;
                    }
                    else
                    {
                        this.Sample = pathIntegrator.Sampler.GetSample();
                        this.Sample.InitSample(SamplesPerVertex * scene.MaxPathDepth);
                    }
                    this.mutationsCount = 0;
                    break;
                case PssmltSamplerState.SmallMutation:
                    mutationsCount++;
                    if (mutationsCount > PssMltEngineController.MaxMutations)
                    {
                        this.SamplerState = PssmltSamplerState.Initialize;
                        goto restart;
                    }
                    if (pathIntegrator.Sampler.GetLazyValue(this.Sample) < PssMltEngineController.LargeStepProbability)
                    {
                        this.SamplerState = PssmltSamplerState.LargeMutation;
                        goto restart;
                    }

                    if (this.XSample == null)
                        this.XSample = new Sample(this.Sample);
                    this.Sample = Mutate(this.Sample, mutationsCount, pathIntegrator.Sampler);
                    break;
            }
            RayData ray;
            if (Sample == null)
            {
                Debugger.Break();
            }
            pathIntegrator.Scene.Camera.GetRay(Sample.imageX, Sample.imageY, out ray);
            this.PathRay = ray;
            this.RayIndex = -1;
            this.depth = 0;
            this.specularBounce = true;
        }


        public override void Splat(SampleBuffer sampleBuffer)
        {
            switch (SamplerState)
            {
                case PssmltSamplerState.LargeMutation:
                case PssmltSamplerState.Initialize:
                    sampleBuffer.SplatSample(this.Sample.imageX, this.Sample.imageY, ref Radiance);
                    if (!Radiance.IsBlack())
                    {
                        this.XContribution = Radiance;
                        this.XSample = new Sample(Sample);
                        this.SamplerState = PssmltSamplerState.SmallMutation;
                        this.mutationsCount = 0;
                    }
                    break;
                case PssmltSamplerState.SmallMutation:
                    if (!Radiance.IsBlack())
                    {
                        var yContrib = Radiance;
                        var yImportance = PssMltEngineController.Importance(ref yContrib);
                        var xImportance = PssMltEngineController.Importance(ref XContribution);

                        //if (!xImportance.NearEqual(yImportance))
                        //{
                        //    Debugger.Break();
                        //}
                        var a = Math.Min(1, (yImportance) / (xImportance));

                        var b_x = contributionMap[XSample].Radiance /Math.Max(1, contributionMap[XSample].Weight);
                        var b_y = contributionMap[this.Sample].Radiance / (Math.Max(1, contributionMap[Sample].Weight));

                        var xPixelContrib = ((1.0 - a) / (PssMltEngineController.MaxMutations)) * (XContribution / (xImportance / b_x));
                        var yPixelContrib = (a / PssMltEngineController.MaxMutations) * (yContrib / (yImportance / b_y));

                        sampleBuffer.SplatSample(XSample.imageX, XSample.imageY, ref xPixelContrib);
                        sampleBuffer.SplatSample(Sample.imageX, Sample.imageY, ref yPixelContrib);

                        contributionMap[XSample].Radiance += xPixelContrib;
                        contributionMap[Sample].Radiance += yPixelContrib;
                        //contributionMap[XSample].Weight += 1.0f;
                        //contributionMap[Sample] .Weight += 1.0f;


                        if (contributionMap.Rndf() < a)
                        {
                            this.XSample = new Sample(Sample);
                            this.XContribution = yContrib;
                        }
                    }
                    break;
                case PssmltSamplerState.WarmUp:
                    this.contributionMap[this.Sample].Radiance += Radiance;
                    this.contributionMap[this.Sample].Weight += 1.0f;
                    sampleBuffer.SplatSample(this.Sample.imageX, this.Sample.imageY, ref Radiance);
                    break;
            }
            this.InitPath(pathIntegrator);
        }

        public override bool FillRayBuffer(RayBuffer rayBuffer)
        {
            var leftSpace = rayBuffer.LeftSpace();
            if (((PathState == PathTracerPathState.EyeVertex) && (1 > leftSpace)) ||
            ((PathState == PathTracerPathState.ShadowRaysOnly) && (tracedShadowRayCount > leftSpace)) ||
            ((PathState == PathTracerPathState.NextVertex) && (tracedShadowRayCount + 1 > leftSpace)))
                return false;
            if (PathState != PathTracerPathState.ShadowRaysOnly)
                RayIndex = rayBuffer.AddRay(ref PathRay);
            if (PathState == PathTracerPathState.NextVertex || PathState == PathTracerPathState.ShadowRaysOnly)
            {
                for (int i = 0; i < tracedShadowRayCount; ++i)
                    secRays[i].ShadowRayIndex = rayBuffer.AddRay(ref secRays[i].ShadowRay);
            }
            return true;
        }

        public override void Advance(RayBuffer rayBuffer, SampleBuffer consumer)
        {
            int currentTriangleIndex = 0;

#if VERBOSE

            try
            {

#endif

                if (((PathState == PathTracerPathState.ShadowRaysOnly) || (PathState == PathTracerPathState.NextVertex)) &&
               (tracedShadowRayCount > 0))
                {
                    var continueRays = new List<ShadowRayInfo>();
                    for (int i = 0; i < tracedShadowRayCount; ++i)
                    {
                        RayHit shadowRayHit = rayBuffer.rayHits[secRays[i].ShadowRayIndex];
                        RgbSpectrum attenuation;
                        bool continueTrace;
                        if (this.ShadowRayTest(ref shadowRayHit, ref secRays[i].ShadowRay, out attenuation, out continueTrace))
                        {
                            Radiance += attenuation * ((secRays[i].Throughput) / secRays[i].Pdf);
                        }

                        if (continueTrace)
                        {
                            continueRays.Add(secRays[i]);
                        }
                    }

                    if (PathState == PathTracerPathState.ShadowRaysOnly)
                    {
                        Splat(consumer);
                        return;
                    }

                    Array.Copy(continueRays.ToArray(), secRays, continueRays.Count);
                    tracedShadowRayCount = continueRays.Count;
                }
                RayHit rayHit = rayBuffer.rayHits[RayIndex];

                depth++;
                Vector wo = -PathRay.Dir;
                bool missed = rayHit.Index == 0xffffffffu;
                if (missed || PathState == PathTracerPathState.ShadowRaysOnly  || depth > scene.MaxPathDepth)
                {
                    if (specularBounce && missed)
                    {
                        Radiance += this.scene.SampleEnvironment(ref wo) * Throughput;
                    }
                    Splat(consumer);

                    return;
                }

                // Something was hit
                if (hitInfo == null)
                {
                    hitInfo = SurfaceSampler.GetIntersection(ref PathRay, ref rayHit);
                }
                else
                {
                    SurfaceSampler.GetIntersection(ref PathRay, ref rayHit, ref hitInfo);
                }

                currentTriangleIndex = (int)rayHit.Index;

                if (hitInfo == null)
                {
                    Debugger.Break();
                }
                //If Hit light)
                if (hitInfo.IsLight)
                {
                    if (specularBounce)
                    {
                        var lt = scene.GetLightByIndex(currentTriangleIndex);
                        if (lt != null)
                        {
                            var le = (RgbSpectrum)(RgbSpectrumInfo)(lt.Le(ref PathRay.Dir));
                            Radiance += Throughput * le;
                            mutate = true;
                        }
                    }
                    Splat(consumer);

                    return;
                }

                var hitPoint = PathRay.Point(rayHit.Distance);

                tracedShadowRayCount = 0;
                var bsdf = hitInfo.MMaterial;

                if (bsdf.IsDiffuse())
                {
                    float lightStrategyPdf = LightSampler.StrategyPdf;
                    //scene.ShadowRaysPerSample/(float)scene.Lights.Length;
                    RgbSpectrum lightTroughtput = Throughput * hitInfo.Color;
                    LightSampler.EvaluateShadow(ref hitPoint, ref hitInfo.Normal, Sample.GetLazyValue(depth * 5),
                                                Sample.GetLazyValue(depth * 6), Sample.GetLazyValue(depth * 7), ref ls);
                    for (int index = 0; index < ls.Length; index++)
                    {

                        if (ls[index].Pdf <= 0f)
                            continue;
                        secRays[tracedShadowRayCount].Throughput = (RgbSpectrum)(RgbSpectrumInfo)(ls[index].Spectrum);
                        secRays[tracedShadowRayCount].Pdf = ls[index].Pdf;
                        secRays[tracedShadowRayCount].ShadowRay = ls[index].LightRay;
                        Vector lwi = secRays[tracedShadowRayCount].ShadowRay.Dir;
                        RgbSpectrum fs;
                        hitInfo.MMaterial.f(ref secRays[tracedShadowRayCount].ShadowRay.Dir, ref wo, ref hitInfo.ShadingNormal, ref Throughput, out fs);
                        secRays[tracedShadowRayCount].Throughput *= lightTroughtput *
                                                                    Vector.AbsDot(ref hitInfo.Normal, ref lwi) *
                                                                    fs;
                        if (!secRays[tracedShadowRayCount].Throughput.IsBlack())
                        {
#if DEBUG
                            RayDen.Library.Core.Components.Assert.IsNotNaN(secRays[tracedShadowRayCount].Pdf);
#endif
                            secRays[tracedShadowRayCount].Pdf /= lightStrategyPdf;
                            tracedShadowRayCount++;
                        }
                    }
                }




                float fPdf;
                Vector wi;

                RgbSpectrum f = hitInfo.MMaterial.Sample_f(ref wo, out wi, ref hitInfo.Normal, ref hitInfo.ShadingNormal, ref Throughput,
                                                          Sample.GetLazyValue(depth * 1), Sample.GetLazyValue(depth * 2),
                                                          Sample.GetLazyValue(depth * 3), ref hitInfo.TextureData,
                                                          out fPdf, out specularBounce);

                if ((fPdf <= 0.0f) || f.IsBlack())
                {
                    if (tracedShadowRayCount > 0)
                        PathState = PathTracerPathState.ShadowRaysOnly;
                    else
                    {
                        Splat(consumer);

                    }
                    return;
                }
                Throughput *= (f * hitInfo.Color) / fPdf;

                if (depth > scene.MaxPathDepth)
                {
                    float prob = Math.Max(Throughput.Filter(), scene.RussianRuletteImportanceCap);
                    if (prob >= Sample.GetLazyValue(depth * 4))
                    {
                        Throughput /= prob;
                    }
                    else
                    {
                        if (tracedShadowRayCount > 0)
                            PathState = PathTracerPathState.ShadowRaysOnly;
                        else
                        {
                            Splat(consumer);

                        }

                        return;
                    }
                }

                PathRay.Org = hitPoint;
                PathRay.Dir = wi.Normalize();
                PathState = PathTracerPathState.NextVertex;


#if VERBOSE
            }
            catch (Exception ex)
            {
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine("Advance path exception");
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine("Error triangle {0}", currentTriangleIndex);
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine(ex.Message);
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine(ex.StackTrace);

            }
#endif

        }

        private Sample Mutate(Sample s, int mc, ISampler rnd)
        {

            if (mc > 1 && rnd.NextFloat() > 0.5f)
            {
                if (rnd.NextFloat() > 0.5f)
                {
                    s.imageX = Mutate(s.imageX, rnd.NextFloat(), rnd.NextFloat());
                    //if (rnd.NextFloat() > 0.5f)
                    s.imageY = Mutate(s.imageY, rnd.NextFloat(), rnd.NextFloat());
                }
                //return s;
                //if (rnd.NextFloat() > 0.5f)return;
            }

            for (int i = 0; i < s.sampleData.Length; i++)
            {
                s.sampleData[i] = Mutate(s.sampleData[i], rnd.NextFloat(), rnd.NextFloat());
            }
            return s;
        }

        private static float Mutate(float value, float u, float v)
        {
            float s1 = 1f / 1024f, s2 = 1f / 64f;
            float dv = (float)(s2 * Math.Exp(-Math.Log(s2 / s1) * u));
            if (v < 0.5)
            {
                value += dv; if (value > 1) value -= 1;
            }
            else
            {
                value -= dv; if (value < 0) value += 1;
            }
            return value;
        }

        private static float Mutate(double value, float u, float v)
        {
            float s1 = 1f / 1024f, s2 = 1f / 64f;
            float dv = (float)(s2 * Math.Exp(-Math.Log(s2 / s1) * u));
            if (v < 0.5)
            {
                value += dv; if (value > 1) value -= 1;
            }
            else
            {
                value -= dv; if (value < 0) value += 1;
            }
            return (float)value;
        }

        protected override bool ShadowRayTest(ref RayHit shadowRayHit, ref RayData shadowRay, out RgbSpectrum attenuation, out bool continueTrace)
        {
            var hit = shadowRayHit.Index == 0xffffffffu || scene.IsLight((int)shadowRayHit.Index);
            attenuation = new RgbSpectrum(1f);
            continueTrace = false;

            if (hit) return true;
            var mesh = scene.GetMeshByTriangleIndex((int)shadowRayHit.Index);
            var mat =
                //SurfaceMaterials.CreateMaterial(scene.MaterialProvider.Get(mesh.MaterialName));
                scene.MatLib.GetSurfMat(mesh != null ? mesh.MaterialName : "default");

            if (mat.TransparentShadows)
            {
                continueTrace = true;
                return true;
            }

            if (mat.AlphaTexture != null)
            {
                continueTrace = true;
                this.SurfaceSampler.GetIntersection(ref shadowRay, ref shadowRayHit, ref hitInfo, IntersectionOptions.ResolveTextures);
                attenuation = (RgbSpectrum)(RgbSpectrumInfo) hitInfo.TextureData.Alpha;
                hit = true;
                //mat.MaterialData.Kt.Filter() < Sample.GetLazyValue();
            }

            return hit;
        }
    }
}
